import os
import pickle

import torch
from torch.utils.data.dataloader import DataLoader

from segmentation.dataset import Handmask, Triplet, MR_slice
from segmentation.utils import *


def get_UNet(
    test_dir,
    folder_dir,
    eval,
    n_train=3,
    batch_size=16,
    split=[72,8,20],
    length=128,
    w_steps=1000,
    combine=True,
    multi_class=True,
):

    import warnings
    warnings.filterwarnings('ignore')

    condition = 0 if 'OASIS' in test_dir else 1

    if eval == 'supr':
        eval_f = eval_dice_3d
        v = f'visualization'
    elif eval == 'ds':
        eval_f = eval_dice_3d_ds
        v = f'visualization_{eval}'
    elif eval == 'myds':
        eval_f = eval_dice_3d_myds
        v = f'visualization_{eval}'
    else:
        print('wrong eval')
        exit()

    if not os.path.exists(os.path.join(folder_dir, v)):
        os.mkdir(os.path.join(folder_dir, v))

    # folder_dir = os.path.join(out_dir, f'exp{id}-{shots}shot-{threshold}thres-{which_net}')
    # init dataaset ===============================================================
    # split 2:1 first, then choose n shot for training, ensure same test set with U-Net training
    split_length = len(os.listdir(os.path.join(test_dir, 'images')))//128

    out_ch = 17
    sum_dice = 0
    sum_dice_class = [0]*(out_ch-1)

    import random
    for i_train in range(n_train):
        torch.manual_seed(i_train)
        random.seed(28*i_train)
        samples = list(range(split_length))

        test_loader = DataLoader(
            MR_slice(
                data_dir=test_dir,
                aug=False,
                sample=samples
            ),
            batch_size=batch_size, drop_last=False, shuffle=False, num_workers=4, pin_memory=True
        )

        print(f'test len: {test_loader.__len__()}')

        chn = 1
        print(f'image channel: {chn}')

        with torch.no_grad():
            with open(os.path.join(folder_dir, 'checkpoint', f'{i_train}_best.pth'), 'rb') as f:
                net = pickle.load(f)['net'].eval().requires_grad_(False).cuda()

            # dice, dice_std, dice_class = eval_dice_3d_ds(net, test_loader, True, folder_dir, i_train, multi_class, True, condition)
            dice, dice_std, dice_class = eval_f(net, test_loader, True, folder_dir, i_train, multi_class, True, condition)

            for d in dice_class:
                print(f'{d:.3f}', end=', ')
            print(f'\ndice: {dice:.3f} std: {dice_std:.3f}\n')

        with open(os.path.join(folder_dir, 'INFO_ds.txt'), 'a') as log:
            log.write(f'mdice: [{dice:.5f}]\n')

        with open(os.path.join(folder_dir, 'test_INFO_ds.txt'), 'a') as log:
            # mdice
            for d in dice_class:
                log.write(f'{d:.3f}, ')
            log.write('\n')
            log.write(f'{dice:.3f}\n')

            sum_dice += dice
            sum_dice_class = [x+y for x,y in zip(sum_dice_class, dice_class)]

            log.write('=================================\n')

    with open(os.path.join(folder_dir, 'INFO_ds.txt'), 'a') as log:
        log.write(f'Mean test miou: [{sum_dice/n_train}]\n')

    with open(os.path.join(folder_dir, 'test_INFO_ds.txt'), 'a') as log:
        log.write(f'Mean {eval}:\n')
        sum_dice_class = [x/n_train for x in sum_dice_class]

        # mdice
        for d in sum_dice_class:
            log.write(f'{d:.3f}, ')
        log.write('\n')
        log.write(f'{sum_dice/n_train:.3f}\n')
        log.write('=================================\n')

if __name__ == "__main__":
    # import time
    # time.sleep(15*60)
    os.environ["CUDA_VISIBLE_DEVICES"] = '7'
    # get_UNet(
    #     test_dir='data/CANDI-128-160-norm',
    #     folder_dir='save_seg/supervised/exp25-357shot-reso160-augTrue-ds',
    #     batch_size=16,
    #     eval='myds',
    #     n_train=3
    # )
    # get_UNet(
    #     test_dir='data/CANDI-128-160-norm',
    #     folder_dir='save_seg/supervised/exp25-357shot-reso160-augTrue-ds',
    #     eval='myds',
    #     n_train=3
    # )
    get_UNet(
        test_dir='data/OASIS-128-160-norm',
        folder_dir='save_seg/supervised/exp35-103shot-reso160-augTrue-ds',
        batch_size=64,
        eval='ds',
        n_train=1
    )
    get_UNet(
        test_dir='data/OASIS-128-160-norm',
        folder_dir='save_seg/supervised/exp35-103shot-reso160-augTrue-ds',
        batch_size=64,
        eval='myds',
        n_train=1
    )
#----------------------------------------------------------------------------
